"""
Memory Management Interface for Continual Learning

This interface defines the standard methods for managing episodic memory
in continual learning scenarios, including exemplar construction and maintenance.
"""

from abc import ABC, abstractmethod
from typing import Tuple, Optional, Any
import numpy as np
import torch


class MemoryManager(ABC):
    """Interface for managing episodic memory in continual learning"""

    @abstractmethod
    def build_rehearsal_memory(self, data_manager: Any, per_class: int) -> None:
        """
        Build or update the rehearsal memory for the current task

        Args:
            data_manager: Data manager providing access to task data
            per_class: Number of exemplars per class to maintain
        """
        pass

    @abstractmethod
    def get_memory(self) -> Optional[Tuple[np.ndarray, np.ndarray]]:
        """
        Retrieve current memory contents

        Returns:
            Tuple of (data, targets) if memory exists, None otherwise
        """
        pass

    @abstractmethod
    def get_exemplar_size(self) -> int:
        """Get the total number of exemplars in memory"""
        pass

    @abstractmethod
    def get_samples_per_class(self) -> int:
        """Get the number of samples per class in memory"""
        pass

    @abstractmethod
    def reduce_exemplar(self, data_manager: Any, m: int) -> None:
        """
        Reduce the number of exemplars per class to m

        Args:
            data_manager: Data manager for accessing data
            m: Target number of exemplars per class
        """
        pass

    @abstractmethod
    def construct_exemplar(self, data_manager: Any, m: int) -> None:
        """
        Construct exemplars for new classes

        Args:
            data_manager: Data manager for accessing data
            m: Number of exemplars per class to construct
        """
        pass


class ExemplarSelector(ABC):
    """Interface for selecting exemplars from class data"""

    @abstractmethod
    def select_exemplars(
        self,
        data: np.ndarray,
        features: np.ndarray,
        num_exemplars: int
    ) -> np.ndarray:
        """
        Select exemplars from class data

        Args:
            data: Raw data samples
            features: Feature representations of the data
            num_exemplars: Number of exemplars to select

        Returns:
            Indices of selected exemplars
        """
        pass


class MemoryStrategy(ABC):
    """Interface for different memory management strategies"""

    @abstractmethod
    def update_memory(
        self,
        new_data: np.ndarray,
        new_targets: np.ndarray,
        current_memory: Optional[Tuple[np.ndarray, np.ndarray]] = None
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Update memory with new data according to strategy

        Args:
            new_data: New data to potentially add to memory
            new_targets: Targets for new data
            current_memory: Current memory state

        Returns:
            Updated memory (data, targets)
        """
        pass

    @abstractmethod
    def get_memory_budget(self, total_classes: int) -> int:
        """
        Calculate memory budget for current number of classes

        Args:
            total_classes: Total number of classes seen so far

        Returns:
            Memory budget (total number of exemplars allowed)
        """
        pass


class ClassMeanManager(ABC):
    """Interface for managing class mean representations"""

    @abstractmethod
    def compute_class_means(self, data_manager: Any) -> torch.Tensor:
        """
        Compute class mean representations

        Args:
            data_manager: Data manager for accessing data

        Returns:
            Tensor of class means [num_classes, feature_dim]
        """
        pass

    @abstractmethod
    def get_class_means(self) -> torch.Tensor:
        """Get current class mean representations"""
        pass

    @abstractmethod
    def update_class_means(self, class_means: torch.Tensor) -> None:
        """Update class mean representations"""
        pass
